# 操作pdf
from datetime import datetime

import fitz  # PyMuPDF
from PIL import Image
import re

# 读取excel文件
import pandas as pd
import os

from paddleocr import PaddleOCR
import logging

# 配置日志
import config


class PDFToJPGProcessor:
    def __init__(self, pdf_path, excel_path):
        self.pdf_path = pdf_path
        self.excel_path = excel_path
        self.excel_df = pd.read_excel(self.excel_path)

        self.paddle_ocr_ = PaddleOCR(use_angle_cls=True, lang="ch")

    def process(self, output_dir, dpi=300):
        """
        将PDF文件按页码分拆成一个个的JPG文件，并对JPG编号。

        :param output_dir: 输出目录
        :param dpi: 输出JPG的分辨率，默认为300 DPI
        """
        # 初始化PDF文档
        doc = fitz.open(self.pdf_path)
        total_pages = len(doc)

        # 遍历PDF的每一页
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)

            # 获取页面的像素表示，指定分辨率
            pix = page.get_pixmap(alpha=False, dpi=dpi)

            # 保存JPG图像
            jpg_filename = os.path.join(output_dir, f'page_{page_num + 1}.jpg')
            print(f"正在转换图片 {page_num + 1}/{total_pages}...")
            pix.save(jpg_filename)

            # 识别并命名
            address, jpg_path, property_id = self.recognising_address(jpg_filename)

            # 如果找到了“不动产单元号”，则重命名jpg文件
            if len(property_id) > 0:

                property_id = property_id[0]
                new_filename = f'{property_id}.jpg'
                new_jpg_path = os.path.join(output_dir, new_filename)
                if os.path.exists(new_jpg_path):
                    logging_info = f"已经存在{new_filename}。请注意：如果不是不同的地址指向同一个不动产单元，则可能是识别错了。"
                    config.print_console_and_write_log(logging_info)
                else:
                    os.rename(jpg_path, new_jpg_path)

                    # 输出进度信息
                    logging_info = f"完成图片重命名：{new_filename}。下一步压缩图片"
                    config.print_console_and_write_log(logging_info)

                    # 压缩图片到450k以内
                    self.compress_image(output_dir + "/" + new_filename, 450)
                    # 输出进度信息
                    logging_info = f"压缩图片完成：{new_filename}"
                    config.print_console_and_write_log(logging_info)


            else:
                logging_info = f"警告: 找不到位置 '{address}' 对应的不动产单元号，跳过文件 page_{page_num + 1}.jpg"
                config.print_console_and_write_log(logging_info)

                # 输出进度信息
                logging_info = f"正在压缩图片：page_{page_num + 1}.jpg"
                config.print_console_and_write_log(logging_info)

                # 压缩图片到450k以内
                self.compress_image(output_dir + "/" + f"page_{page_num + 1}.jpg", 450)
                # 输出进度信息
                logging_info = f"压缩图片完成：page_{page_num + 1}.jpg"
                config.print_console_and_write_log(logging_info)

            # 释放Pixmap对象资源
            pix = None

            # 关闭PDF文档
        doc.close()
        print("图片转换完成。")

    def rename_image(self, jpg_path, output_dir, property_id):
        property_id = property_id[0]
        new_filename = f'{property_id}.jpg'
        new_jpg_path = os.path.join(output_dir, new_filename)

        if os.path.exists(new_jpg_path):
            logging_info = f"已经存在{new_filename}"
            config.print_console_and_write_log(logging_info)
        else:
            os.rename(jpg_path, new_jpg_path)
        return new_filename, property_id

    def ocr_extract_text(self, image_path):

        ocr = PaddleOCR(use_angle_cls=True, lang="ch")

        result = ocr.ocr(image_path, cls=True)

        for idx in range(len(result)):
            res = result[idx]
            # print(f"提取的信息为：{res}")
            for index, line in enumerate(res, start=0):
                if line[1][0] == '座':
                    print(f"地址是：{res[index + 1][1][0]}")
                    return res[index + 1][1][0]

    def ocr_extract_text_by_regular_expression(self, image_path):
        result = self.paddle_ocr_.ocr(image_path, cls=True)

        for idx in range(len(result)):
            res = result[idx]
            # print(f"提取的信息为：{res}")
            for line in res:
                target = line[1][0]
                # 检查一个变量是否是特定类型通常使用 isinstance() 函数
                if isinstance(target, str):
                    if target.__contains__('市') or target.__contains__("南沙区"):
                        return target
        logging_info = f"可能识别不出来。整个识别结果为：{res}"
        config.print_console_and_write_log(logging_info)

    def get_location_from_filename(self, filename):
        # 假设文件名格式为 'location_page_X.jpg'，我们需要提取出'location'部分
        parts = filename.split('_')
        if len(parts) >= 2 and parts[0].endswith('.jpg'):
            # 去掉文件扩展名，并假设位置信息是文件名中的第一部分
            location = parts[0][:-4]
            return location
        else:
            # 如果文件名格式不匹配，返回None或抛出异常
            return None

    def replace_O_with_0_in_room_numbers(self, address):
        # 正则表达式模式，匹配 "O" 后面跟着一个或多个数字，然后是 "房"
        pattern = r'O(\d+)房'

        # 替换匹配到的模式，将 "O" 替换为 "0"
        replaced_address = re.sub(pattern, lambda match: '0' + match.group(1) + '房', address)

        return replaced_address

    def check_and_correct_address(self, target_address):
        # 修正城市的简称

        if target_address.startswith('州市'):
            target_address = '广' + target_address
        # 替换错误的街道名分隔符
        target_address = target_address.replace('-街', '一街')

        # 替换错误的符号
        target_address = target_address.replace('—层', '一层')

        # 替换错误的分隔符
        target_address = target_address.replace('-层', '一层')

        # 替换错误的分隔符
        target_address = target_address.replace('-楼', '一楼')
        # 替换错误的分隔符
        target_address = target_address.replace('-户', '一户')

        # O替换0
        target_address = target_address.replace('O', '0')

        # 街号
        target_address = target_address.replace('街号', '街1号')

        target_address = self.replace_O_with_0_in_room_numbers(target_address)

        # 如果有一个字符串中包括有XX房，就把房之后的字符串舍弃
        if '房' in target_address:
            target_address = target_address[:target_address.find('房') + 1]

        # 如果有一个字符串中包括有单元，就把单元之后的字符串舍弃
        unit_index = target_address.find("单元")
        if unit_index != -1:
            target_address = target_address[:unit_index + 2]  # +2 是为了包括“单元”这两个字符

        # 如果有一个字符串中包括有XX车位，就把车位之后的字符串舍弃
        if '车位' in target_address:
            target_address = target_address[:target_address.find('车位') + 1]

        # 如果以“车”结尾，则补全
        if target_address.endswith('车'):
            target_address = target_address + '位'

        # 删除数字与房之间的空格
        target_address = re.sub(r'\s+房', '房', target_address)

        # 查找"图幅号"的位置
        index = target_address.find('图幅号')
        # 如果找到了"图幅号"
        if target_address.find('图幅号') != -1:
            # 使用rsplit从右边开始分割字符串，移除"图幅号"及其之后的部分
            # maxsplit参数确保只分割第一次出现的位置
            target_address = target_address.rsplit('图幅号', 1)[0]
        # 如果没有找到"图幅号"，则返回原字符串

        return target_address

    def compress_image(self, image_path, max_size_kb):
        """
        压缩图片到指定大小以下。

        :param image_path: 图片的路径
        :param max_size_kb: 目标最大大小（以KB为单位）
        :return: 压缩后的图片路径
        """
        # 打开图片
        img = Image.open(image_path)

        # 获取图片原始大小
        original_size = os.path.getsize(image_path)

        # 如果图片已经小于目标大小，则直接返回
        if original_size <= max_size_kb * 1024:
            return image_path

            # 设置最大文件大小为KB转为字节
        max_size = max_size_kb * 1024

        # 转换图片为RGB模式
        img = img.convert('RGB')

        # 设置初始质量为95
        quality = 95

        # 尝试降低图片质量来压缩图片
        while quality >= 0:
            # 保存图片并设置质量参数
            img.save(image_path, 'JPEG', quality=quality)

            # 获取压缩后图片的大小
            new_size = os.path.getsize(image_path)

            # 如果图片大小小于目标大小，则停止压缩
            if new_size <= max_size:
                break

                # 否则，继续降低图片质量
            quality -= 5  # 降低质量参数

        # 如果质量降到0仍然无法满足大小要求，则抛出异常
        if quality < 0:
            raise ValueError("Could not compress image to desired size.")

        return image_path

    def batch_compress_images(self, folder_path, max_size_kb):
        """
        批量压缩文件夹中的所有JPEG图片到指定大小以下。

        :param folder_path: 包含图片的文件夹路径
        :param max_size_kb: 目标最大大小（以KB为单位）
        """
        # 遍历文件夹中的所有文件
        for filename in os.listdir(folder_path):
            # 检查文件是否为图片（这里只检查.jpg和.jpeg文件）
            if filename.lower().endswith(('.jpg', '.jpeg')):
                # 构建完整的文件路径
                image_path = os.path.join(folder_path, filename)

                # 压缩图片
                try:
                    compressed_image_path = self.compress_image(image_path, max_size_kb)
                    print(f'Compressed {image_path} to {compressed_image_path}')
                except ValueError as e:
                    print(f'Could not compress {image_path}: {e}')

    def rename_jpg_with_property_id(self, jpg_dir):
        # 获取所有JPG文件的文件名
        jpg_filenames = [f for f in os.listdir(jpg_dir) if f.endswith('.jpg')]
        # 遍历jpg文件夹中的所有文件
        for index, filename in enumerate(jpg_filenames, start=1):
            if filename.endswith('.jpg'):
                address, jpg_path, property_id = self.recognising_address(jpg_dir + filename)

                self.rename_image_with_bdcdyh(address, filename, index, jpg_dir, jpg_filenames, jpg_path, property_id)

        print("图片重命名完成。")

    def rename_image_with_bdcdyh(self, address, filename, index, jpg_dir, jpg_filenames, jpg_path, property_id):
        # 如果找到了“不动产单元号”，则重命名jpg文件
        if len(property_id) > 0:
            new_filename, property_id = self.rename_image(jpg_path, jpg_dir, property_id)

            # 输出进度信息
            print(f"正在重命名图片 {index}/{len(jpg_filenames)}: {filename} -> {new_filename}")
            with open(config.log_file_path, 'a') as log_file:
                log_file.write(f"正在重命名图片 {index}/{len(jpg_filenames)}: {filename} -> {new_filename}\n")
        else:
            print(f"警告: 找不到位置 '{address}' 对应的不动产单元号，跳过文件 {filename}")
            with open(config.log_file_path, 'a') as log_file:
                log_file.write(f"警告: 找不到位置 '{address}' 对应的不动产单元号，跳过文件 {filename}\n")

    def recognising_address(self, image_path):
        print(f"分户图路径为：" + image_path)
        # 提取文字
        address = self.ocr_extract_text_by_regular_expression(image_path)

        property_id = ""
        if address is None:
            logging_info = f"当前识别结果为空！"
            config.print_console_and_write_log(logging_info)
        else:

            address = self.check_and_correct_address(address)
            logging_info = f"分户图提取的‘座落’为：" + address
            config.print_console_and_write_log(logging_info)

            # 在Excel表中查找不动产单元号
            property_id = self.excel_df.loc[self.excel_df['zl'] == address, 'bdcdyh'].values
            logging_info = f"从关系表中查找到的不动产单元号为：{property_id}"
            config.print_console_and_write_log(logging_info)

        return address, image_path, property_id
